# common/MQTTClientWrapper.py

import network
import time
from umqtt.simple import MQTTClient

class MQTTClientWrapper:
    def __init__(self, ssid, password, server, port, client_id, topic_sub, topic_pub):
        self.ssid = ssid
        self.password = password
        self.server = server
        self.port = port
        self.client_id = client_id
        self.topic_sub = topic_sub
        self.topic_pub = topic_pub
        self.client = None
        self.latest_message = None

    def connect_wifi(self):
        wlan = network.WLAN(network.STA_IF)
        wlan.active(True)
        if not wlan.isconnected():
            print('Connecting to network...')
            wlan.connect(self.ssid, self.password)
            while not wlan.isconnected():
                pass
        print('Network config:', wlan.ifconfig())
        return wlan

    def connect_mqtt(self):
        self.client = MQTTClient(self.client_id, self.server, port=self.port, keepalive=60)
        try:
            self.client.connect()
            print('Connected to %s MQTT broker' % self.server)
            print('连接MQTT成功')
            self.client.set_callback(self.sub_callback)
            self.client.subscribe(self.topic_sub)
            print('Subscribed to topic: %s' % self.topic_sub)
        except Exception as e:
            print('连接MQTT失败: %s' % e)
            self.client = None

    def sub_callback(self, topic, msg):
        self.latest_message = (topic, msg)
        print(f"Received message: {msg} from topic: {topic}")
        print(f"Latest message updated: {self.latest_message}")

    def check_msg(self):
        if not self.client:
            return
        try:
            self.client.check_msg()
        except OSError as e:
            print("Error in check_msg:", e)
            self.reconnect_mqtt()

    def reconnect_mqtt(self):
        try:
            self.client.disconnect()
        except:
            pass
        time.sleep(5)  # 等待一段时间再尝试重新连接
        try:
            self.client.connect()
            print('Reconnected to %s MQTT broker' % self.server)
            self.client.set_callback(self.sub_callback)
            self.client.subscribe(self.topic_sub)
        except Exception as e:
            print("Failed to reconnect to MQTT broker:", e)

    def get_latest_message(self):
        return self.latest_message

    def publish(self, message):
        if self.client:
            self.client.publish(self.topic_pub, message)
            print('Published message to topic: %s' % self.topic_pub)

    def start(self):
        wlan = self.connect_wifi()
        if not wlan.isconnected():
            print("WiFi connection failed, exiting.")
            return False
        self.connect_mqtt()
        if not self.client:
            print("MQTT connection failed, exiting.")
            return False
        return True